import logging
import os 
from gensim.models import word2vec

def train(path):
    logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s',level=logging.INFO)
    sentences = word2vec.LineSentence(path)
    model = word2vec.Word2Vec(sentences,hs=1,min_count=5,window=5,size=100)
    return model

def test(model,test_name):
    # req_count = 5
    # for key in model.wv.similar_by_word(test_name, topn =100):
    #     req_count -= 1
    #     print(key[0], key[1])
    #     if req_count == 0:
    #         break

    req_count = 5
    for key in model.most_similar(positive=['爱国路','泰宁花园'], negative=['城市天地广场'], topn =100):
        req_count -= 1
        print(key[0], key[1])
        if req_count == 0:
            break

if __name__ == "__main__":
    # model = train('word2vec/raw_result.txt')
    # model.save("word2vec/mymodel")

    model = word2vec.Word2Vec.load("word2vec/mymodel")
    test(model,'景田东路')

